import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow.contrib.slim.python.slim.nets import resnet_v1
pretrainedModelDir = "/home/dolly/checkpoints/resnet_v1_101_2016_08_28/resnet_v1_101.ckpt"
outputPath = "/home/dolly/checkpoints/resnet_v1_101_2016_08_28/resnet_v1_101.pb"

height, width = 224, 224
X = tf.placeholder(tf.float32, [None, height, width, 3])  # 定义输入张量
Y = tf.placeholder(tf.float32, [None, 10])
# 读取网络
with slim.arg_scope(resnet_v1.resnet_arg_scope()):
    logits, end_points = resnet_v1.resnet_v1_101(X, num_classes=1000, is_training=False)


with tf.Session() as sess:
    # 先初始化所有变量，避免有些变量未读取而产生错误
    init = tf.global_variables_initializer()
    sess.run(init)

    # 这里的exclusions是不需要读取预训练模型中的Logits,因为默认的类别数目是1000，当你的类别数目不是1000的时候，如果还要读取的话，就会报错
    # exclusions = ['InceptionV3/Logits',
    #               'InceptionV3/AuxLogits']
    exclusions = []
    # 创建一个列表，包含除了exclusions之外所有需要读取的变量
    inception_except_logits = slim.get_variables_to_restore(exclude=exclusions)

    # 建立一个从预训练模型checkpoint中读取上述列表中的相应变量的参数的函数
    init_fn = slim.assign_from_checkpoint_fn(pretrainedModelDir, inception_except_logits, ignore_missing_vars=False)
    # 将权重加载到sess中
    init_fn(sess)
    # print('Loaded.')

    tensor_name_list = [tensor.name for tensor in tf.get_default_graph().as_graph_def().node]
    for tensor_name in tensor_name_list:
        print("node name: %s\n" % tensor_name)
